import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0, 1'

import csv
import pickle
import random

import numpy as np
import torch


if __name__ == "__main__":
    input_path = "/home/ubuntu/Ascend/model_inference/pytorch_source/crop_inputs/torch.add/torch.add.pickle"
    compare_result_path = "/home/ubuntu/Ascend/model_inference/pytorch_results/crop_0d/torch.add/20230327_121459/torch.add.csv"
    source_result_path = "/home/ubuntu/Ascend/model_inference/pytorch_results/crop_0d/torch.add/20230327_121459/torch.add_source.pickle"
    follow_result_path = "/home/ubuntu/Ascend/model_inference/pytorch_results/crop_0d/torch.add/20230327_121459/torch.add_follow.pickle"

    with open(input_path, "rb") as f:
        input_dict_list = pickle.load(f)

    source_results = []
    follow_results = []
    compare_results = []

    cnt = 0
    for input_dict in input_dict_list:
        for key in input_dict.keys():
            val = input_dict[key]
            if type(val) is np.ndarray:
                input_dict[key] = torch.from_numpy(val).cuda()
        crop_input = torch.tensor([])
        crop_key = ""
        for key in input_dict.keys():
            val = input_dict[key]
            if type(val) is torch.Tensor and crop_input.numel() < val.numel():
                crop_input = val
                crop_key = key
        if crop_key == "" or crop_input.numel() == 0 or crop_input.shape[0] < 2:
            source_results.append("Pass")
            follow_results.append("Pass")
            compare_results.append([cnt, "Pass"])
            cnt += 1
            continue
        axis_0d = crop_input.shape[0]
        random_0d = random.randint(1, axis_0d - 1)
        pos_0d = [random_0d, axis_0d - random_0d]
        random_0d = random.randint(0, 1)
        if random_0d == 0:
            follow_input_0d = torch.split(crop_input, dim=0, split_size_or_sections=pos_0d)[0]
        else:
            follow_input_0d = torch.split(crop_input, dim=0, split_size_or_sections=pos_0d)[1]
        a = "Error"
        try:
            a = torch.add(**input_dict)
            print("Full output:")
            print(a)
            a = a.cpu().numpy()
            source_results.append(a)
        except Exception as e:
            print("[Crop] Exception occurred~:\n", e)
            source_results.append("Error")

        input_dict[crop_key] = follow_input_0d
        try:
            b = torch.add(**input_dict)
            print("Rule 5 output:")
            print(b)
            b = b.cpu().numpy()
            if a == "Error":
                compare_results.append([cnt, "Source Error"])
            elif type(a) is np.ndarray and type(b) is np.ndarray:
                temp = torch.split(torch.from_numpy(a).cuda(), dim=0, split_size_or_sections=pos_0d)[random_0d].cpu().numpy()
                if temp.size != b.size:
                    compare_results.append([cnt, "Shape unmatch", pos_0d, random_0d])
                elif np.allclose(temp, b, atol=1.e-2, rtol=1.e-5, equal_nan=True):
                    compare_results.append([cnt, "Consistent", a.shape, b.shape, a.dtype, b.dtype])
                else:
                    compare_results.append([cnt, "Inconsistent", a.shape, b.shape, a.dtype, b.dtype, pos_0d, random_0d])
            else:
                compare_results.append([cnt, "Not Tensor"])
            follow_results.append(b)
        except Exception as e:
            print("[Crop] Exception occurred~:\n", e)
            follow_results.append("Error")
            compare_results.append([cnt, "Follow Error"])

        cnt += 1

    with open(source_result_path, "wb+") as f:
        pickle.dump(source_results, f)

    with open(follow_result_path, "wb+") as f:
        pickle.dump(follow_results, f)

    with open(compare_result_path, "w+", newline="") as f:
        w = csv.writer(f)
        w.writerows(compare_results)
